import os

from dotenv import load_dotenv
from openai import OpenAI

from prompt import propose_prompt, value_prompt

load_dotenv()
api_key = os.getenv("api_key")
base_url = os.getenv("base_url")
client = OpenAI(api_key=api_key, base_url=base_url)


def chatgpt(prompt, model="gpt-3.5-turbo", n=1) -> list:
    messages = [{"role": "user", "content": prompt}]
    outputs = []
    while n > 0:
        cnt = min(n, 20)
        n -= cnt
        res = client.chat.completions.create(model=model, messages=messages, n=cnt)
        outputs.extend([choice.message.content for choice in res.choices])
    return outputs


# input = input("输入：")
input = "5 8 11 13"

#######################第一轮：获取候选项#############################
print("第一轮：思考生成中......")
# 获得可能的下一步行动：
# str: 2 + 8 = 10 (left: 8 10 14)\n8 / 2 = 4 (left: 4 8 14)
# array:[2 + 8 = 10 (left: 8 10 14),8 / 2 = 4 (left: 4 8 14)]
proposals = chatgpt(propose_prompt.format(input=input))[0].split('\n')
# 每个元素后加上/n （为了拼接后续步骤）
# [2 + 8 = 10 (left: 8 10 14)\n,8 / 2 = 4 (left: 4 8 14)\n]
proposals = [_ + '\n' for _ in proposals]
# 创建对应的index list
ids = list(range(len(proposals)))

print("第一轮思考完毕，结果：")
for item in proposals:
    item = item.replace("\n", "\\n")
    print(item)

# 分析能到达24的可能性
print("第一轮：评估结果......")
values = []
value_cache = []
for proposal in proposals:
    # 2 + 8 = 10 (left: 8 10 14)\n  取出left中的数组  8 10 14
    current_numbers = proposal.strip().split('\n')[-1].split('left: ')[-1].split(')')[0]  # todo 应该加上缓存，去掉可能出现的重复评估项
    # 生成3次： 8 10 14
    # 8+10+14 (left 0.001)\n
    # 8-10+14 (left 0.001)\n
    # 8*10-14 (left 0.001)
    if current_numbers in value_cache:
        print(f"{current_numbers}已评估")
        values.append(0)
        continue
    print(f"评估:{current_numbers}...")
    value_outputs = chatgpt(value_prompt.format(input=current_numbers), n=3)
    value_cache.append(current_numbers)
    # 从三次的结果中取出分值：0.003
    value = 0
    try:
        value = sum([float(_.split('\n')[-1].split('left: ')[-1]) for _ in value_outputs])
    except:
        pass
    values.append(value)
    # 日志：
    temp = proposal.replace("\n", "\\n")
    print(f"[{temp}] 得分: {value}")

print("第一轮：筛选候选项中......")
# 根据预测结果选择前5个最有可能到达24的候选项
# 按照评级排序下一步动作的index 并记录对应的下一步动作
select_ids = sorted(ids, key=lambda i: values[i], reverse=True)[:5]
select_proposals = [proposals[select_id] for select_id in select_ids]

print("第一轮候选项：", select_proposals)

print("####################################################")

####################### 第二轮：从候选项中进一步计算第三个数 #############################
print("第二轮：思考生成中......")
values = []
# 获得可能的下一步动作
proposals = []
for proposal in select_proposals:
    # 从第一轮候选结果中取出数组 获取下一步
    # 2 * 3 = 6 (left: 2 2 6)\n -》 2 2 6
    current_numbers = proposal.strip().split('\n')[-1].split('left: ')[-1].split(')')[0]
    # 2 + 2 = 4 (left: 4 6)\n
    # 2 * 2 = 4 (left: 4 6)\n
    # 6 - 2 = 4 (left: 2 4)\n
    p = chatgpt(propose_prompt.format(input=current_numbers))[0].split('\n')
    # 将第一步的结果和本次拼接
    # ['2 * 3 = 6 (left: 2 2 6)\n2 + 2 = 4 (left: 4 6)\n','2 * 3 = 6 (left: 2 2 6)\n6 - 2 = 4 (left: 2 4)\n']
    for item in p:
        proposals.append(proposal + item + '\n')

print("第二轮思考完毕，结果：")
for item in proposals:
    item = item.replace("\n", "\\n")
    print(item)

ids = list(range(len(proposals)))

print("第二轮：评估结果......")
for proposal in proposals:
    current_numbers = proposal.strip().split('\n')[-1].split('left: ')[-1].split(')')[0]
    if current_numbers in value_cache:
        print(f"{current_numbers}已评估")
        values.append(0)
        continue
    print(f"评估:{current_numbers}...")
    value_outputs = chatgpt(value_prompt.format(input=current_numbers), n=3)
    value_cache.append(current_numbers)
    value = 0
    try:
        value = sum([float(_.split('\n')[-1].split('left: ')[-1]) for _ in value_outputs])
    except:
        pass
    values.append(value)

print("第二轮：筛选候选项中......")
select_ids = sorted(ids, key=lambda i: values[i], reverse=True)[:5]
select_proposals = [proposals[select_id] for select_id in select_ids]

print("第二轮候选项：", select_proposals)

print("####################################################")

####################### 第三轮： #############################
print("第三轮：思考生成中......")
values = []
# 获得可能的下一步动作
proposals = []
for proposal in select_proposals:
    current_numbers = proposal.strip().split('\n')[-1].split('left: ')[-1].split(')')[0]
    p = chatgpt(propose_prompt.format(input=current_numbers))[0].split('\n')
    for item in p:
        proposals.append(proposal + item + '\n')

print("第三轮思考完毕，结果：")
for item in proposals:
    item = item.replace("\n", "\\n")
    print(item)

ids = list(range(len(proposals)))

result = []
print("第三轮：评估结果......")
for proposal in proposals:
    current_numbers = proposal.strip().split('\n')[-1].split('left: ')[-1].split(')')[0]
    if current_numbers == '24':
        result.append(proposal)

print(f"输入的数字为：{input}。")
if len(result) == 0:
    print("无法获得24")
else:
    print("获得24的算式有:")
    for r in result:
        print("=================")
        print(r)
